from dcprogs import read_idealized_bursts
from dcprogs.likelihood import QMatrix

name   = "CH82.scn"
tau    = 1e-4
tcrit  = 4e-3
graph  = [["V", "V", "V",   0,   0],
          ["V", "V",   0, "V",   0],
          ["V",   0, "V", "V", "V"],
          [  0, "V", "V", "V",   0],
          [  0,   0, "V",   0, "V"]]
nopen  = 2
qmatrix = QMatrix([[ -3050,        50,  3000,      0,    0 ],
                  [ 2./3., -1502./3.,     0,    500,    0 ],
                  [    15,         0, -2065,     50, 2000 ],
                  [     0,     15000,  4000, -19000,    0 ],
                  [     0,         0,    10,      0,  -10 ] ], 2)

bursts = read_idealized_bursts(name, tau=tau, tcrit=tcrit)

from scipy.optimize import minimize
from numpy import NaN, zeros, arange
import numpy as np
from dcprogs.likelihood.random import qmatrix as random_qmatrix
from dcprogs.likelihood import QMatrix, Log10Likelihood
from dcprogs.likelihood.optimization import reduce_likelihood

likelihood = Log10Likelihood(bursts, nopen, tau, tcrit)
reduced = reduce_likelihood(likelihood, graph)
x = reduced.to_reduced_coords( random_qmatrix(5).matrix )

constraints = []
def create_inequality_constraints(i, value=0e0, sign=1e0):
    f = lambda x: sign * (x[i]  - value)
    def df(x):
        a = zeros(x.shape)
        a[i] = sign
        return a
    return f, df

for i in range(len(x)):
    f, df = create_inequality_constraints(i)
    constraints.append({'type': 'ineq', 'fun': f, 'jac': df})
    f, df = create_inequality_constraints(i, 1e4, -1)
    constraints.append({'type': 'ineq', 'fun': f, 'jac': df})


def random_starting_point():
    from numpy import infty, NaN
    from dcprogs.likelihood.random import rate_matrix as random_rate_matrix


    for i in range(100):
        matrix = random_rate_matrix(N=len(qmatrix.matrix), zeroprob=0)
        x = reduced.to_reduced_coords( matrix )
        try:
            result = reduced(x)
            print(result, reduced.to_full_coords(x))
        except:
            pass
        else:
            if result != NaN and result != infty and result != -infty: break
    else: raise RuntimeError("Could not create random matrix")
    return x

def does_not_throw(x):
    try: return -reduced(x)
    except: return NaN

import math
methods = ['COBYLA', 'SLSQP']
x = random_starting_point()
print ('x=', x)
maxx = (x.copy(), reduced(x))
for i in range(len(methods)):
    result = minimize(does_not_throw,
                      x,
                      method=methods[i],
                      constraints=constraints,
                      options={'maxiter': 1000, 'disp':True})

    print(result)
    if not math.isnan(result.fun):
        if result.fun < maxx[1]: maxx = (x.copy(), result.fun)
        if result.success and i > 4: break
    x +=  random_starting_point() * 1e-2
    if np.all(np.isnan(x)): x = random_starting_point()
print(maxx[0])
print(maxx[1])